package com.test.html;

import com.alibaba.fastjson.JSONArray;
import com.alibaba.fastjson.JSONObject;
import okhttp3.*;
import okio.BufferedSink;
import okio.Okio;
import okio.Sink;

import java.io.File;
import java.io.FileWriter;
import java.io.IOException;
import java.nio.file.Files;
import java.nio.file.Path;
import java.nio.file.Paths;
import java.time.LocalDate;
import java.time.format.DateTimeFormatter;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * 下载midjourney图片
 */
public class HtmlParser {
    private static AtomicInteger handled = new AtomicInteger(0);

    private static LocalDate now = LocalDate.now();
    private static String date = now.format(DateTimeFormatter.ofPattern("yy-MM"));
    private static String baseDir = "D:\\tmp\\" + date;
    private static String fileName = "jsonData.txt";

    public static void main(String[] args) throws IOException, InterruptedException {
        System.out.println(baseDir);
        Path path = Paths.get(baseDir);
        Files.createDirectories(path);

        /* 下载单个图片
        download("https://cdn.midjourney.com/1b9f6d82-a5dd-4600-98d8-3a54c637a7bd/0_1.png");
        try {
            Thread.sleep(4000L);
        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        }*/

        // 从midjourney下载图片
        String url = // "https://www.midjourney.com/showcase/top/"; //
                "https://www.midjourney.com/showcase/recent/";
        String doc = OkHttpUtilsLocal.builder()
                .url(url)
                .get()
                .sync();
        System.out.println(doc);
        String json = doc.substring(doc.indexOf("<script id=\"__NEXT_DATA__\" type=\"application/json\">") + 51);
        json = json.substring(0, json.indexOf("</script>"));
        System.out.println(json);
        saveToFile(baseDir + "_" + fileName, now.format(DateTimeFormatter.ofPattern("yy-MM-dd")));
        JSONObject data = JSONObject.parseObject(json);
        JSONObject props = data.getJSONObject("props");
        JSONObject pageProps = props.getJSONObject("pageProps");
        JSONArray jobs = pageProps.getJSONArray("jobs");
        int size = jobs.size();
        JSONObject job;
        String imageUrl;
        for (int i = 0; i < size; i++) {
            job = jobs.getJSONObject(i);
            imageUrl = job.getJSONObject("event").getString("seedImageURL");
            File dest = new File(baseDir, getFileNameFromUrl(imageUrl));
            if (dest.exists()) {
                handled.incrementAndGet();
                System.out.println(imageUrl + "已存在");
                continue;
            }
            saveToFile(baseDir + "_" + fileName, job.toJSONString());
            download(imageUrl);
            Thread.sleep(500L);
        }
        System.out.println("已有文件数：" + handled.get());
        while (handled.get() < size) {
            try {
                Thread.sleep(1000L);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }
        System.out.println(handled.get());
        System.exit(1);
    }

    public static void saveToFile(String fileName, String txt) {
        try {
            FileWriter writer = new FileWriter(fileName, true); // 创建一个 FileWriter 对象来写入文件
            writer.write(txt);  // 写入文本内容
            writer.write("\r\n");
            writer.close();   // 关闭 FileWriter 对象
            System.out.println("文件已成功保存。");
        } catch (IOException e) {
            System.out.println("文件保存时出错。");
            e.printStackTrace();
        }
    }

    private static String getFileNameFromUrl(String url) {
        return url.substring(27).replace("/", "_");
    }

    private static final OkHttpClient client = new OkHttpClient();
    private static final int REQUEST_EXTERNAL_STORAGE = 1;

    //下载文件
    private static void downloadFile(final String url) {
        new Thread(() -> download(url)).start();
    }

    //下载
    public static void download(final String url) {

        // System.out.println("DOWNLOAD" + "startTime=" + startTime);
        Request request = new Request.Builder().url(url).build();
        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(Call call, IOException e) {
                handled.incrementAndGet();
                // 下载失败
                e.printStackTrace();
                System.out.println("DOWNLOAD" + "download failed");
            }

            @Override
            public void onResponse(Call call, Response response) throws IOException {
                Sink sink = null;
                BufferedSink bufferedSink = null;
                try {
                    final long startTime = System.currentTimeMillis();
                    File dest = new File(baseDir, getFileNameFromUrl(url)); // url.substring(27, url.lastIndexOf("/")) + ".png"
                    if (dest.exists()) {
                        System.out.println(url + "已存在");
                        return;
                    }
                    sink = Okio.sink(dest);
                    bufferedSink = Okio.buffer(sink);
                    bufferedSink.writeAll(response.body().source());

                    bufferedSink.close();
                    System.out.println(url + " download success " + "totalTime=" + (System.currentTimeMillis() - startTime));
                } catch (Exception e) {
                    e.printStackTrace();
                    System.out.println("DOWNLOAD" + "download failed");
                } finally {
                    handled.incrementAndGet();
                    if (bufferedSink != null) {
                        bufferedSink.close();
                    }
                    response.close();
                }
            }
        });
    }
}
